package com.gitee.wsl.ext.collection

fun <T> Iterable<T>.groupConsecutiveBy(groupIdentifier: (T, T) -> Boolean) =
    if (!this.any())
        emptyList()
    else this
        .drop(1)
        .fold(mutableListOf(mutableListOf(this.first()))) { groups, t ->
            groups.last().apply {
                if (groupIdentifier.invoke(last(), t)) {
                    add(t)
                } else {
                    groups.add(mutableListOf(t))
                }
            }
            groups
        }


fun <T: Any, K: Any, V: Any> Iterable<T>.groupAdjacent(f: (T) -> Pair<K, V>): List<Pair<K, List<V>>> {
    val res = mutableListOf<Pair<K, List<V>>>()
    val group = mutableListOf<V>()
    var groupKey: K? = null

    for (item in this) {
        val (key, value) = f(item)
        if (group.isNotEmpty() && key != groupKey && groupKey != null) {
            res.add(groupKey to group)
            group.clear()
        }
        group.add(value)
        groupKey = key
    }

    if (group.isNotEmpty() && groupKey != null) {
        res.add(groupKey to group)
    }

    return res
}

public inline fun <T, K : Any> Iterable<T>.groupByNotNull(keySelector: (T) -> K?): Map<K, List<T>> =
    mapNotNull { (keySelector(it) ?: return@mapNotNull null) to it }
        .groupBy({ it.first }, { it.second })

public inline fun <T, K : Any, V : Any> Iterable<T>.groupByNotNull(
    keySelector: (T) -> K?,
    valueTransform: (T) -> V?
): Map<K, List<V>> =
    mapNotNull {
        val k = keySelector(it) ?: return@mapNotNull null
        val v = valueTransform(it) ?: return@mapNotNull null
        k to v
    }.groupBy({ it.first }, { it.second })
